@echo off
REM 佛山大学知识问答系统 - Windows启动脚本

setlocal enabledelayedexpansion

REM 颜色定义
set "RED=[31m"
set "GREEN=[32m"
set "YELLOW=[33m"
set "BLUE=[34m"
set "NC=[0m"

REM 日志函数
:log_info
echo %BLUE%[INFO]%NC% %~1
exit /b

:log_success
echo %GREEN%[SUCCESS]%NC% %~1
exit /b

:log_warning
echo %YELLOW%[WARNING]%NC% %~1
exit /b

:log_error
echo %RED%[ERROR]%NC% %~1
exit /b

REM 检查虚拟环境
:check_venv
if not exist ".venv" (
    call :log_error "虚拟环境不存在，请先创建虚拟环境"
    exit /b 1
)
exit /b 0

REM 激活虚拟环境
:activate_venv
call :log_info "激活虚拟环境..."
call .venv\Scripts\activate.bat
exit /b

REM 检查CUDA环境
:check_cuda
call :log_info "检查CUDA环境..."
python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}'); print(f'GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"No GPU\"}')"
exit /b

REM 检查模型文件
:check_models
call :log_info "检查模型文件..."

if not exist "models\Qwen3-4B-Instruct-2507" (
    call :log_warning "Qwen3-4B模型未找到，请先下载模型"
    set /p choice="是否现在下载？(y/n): "
    if /i "!choice!"=="y" (
        call :download_models
    ) else (
        call :log_error "模型文件缺失，无法继续"
        exit /b 1
    )
)

if not exist "models\m3e-base" (
    call :log_warning "m3e-base embedding模型未找到"
    set /p choice="是否现在下载？(y/n): "
    if /i "!choice!"=="y" (
        call :download_embedding_model
    )
)
exit /b

REM 下载模型
:download_models
call :log_info "下载Qwen3-4B模型..."
huggingface-cli download Qwen/Qwen3-4B-Instruct-2507 --local-dir "./models/Qwen3-4B-Instruct-2507" --resume-download
exit /b

:download_embedding_model
call :log_info "下载embedding模型..."
huggingface-cli download moka-ai/m3e-base --local-dir "./models/m3e-base" --resume-download
exit /b

REM 检查数据
:check_data
call :log_info "检查数据文件..."

if not exist "processed_data" (
    call :log_info "处理训练数据..."
    python src/data_processor.py
)

if not exist "university_knowledge" (
    call :log_info "构建RAG向量数据库..."
    python src/rag_builder.py
)
exit /b

REM 训练模型
:train_model
call :log_info "开始LoRA微调训练..."
call :log_warning "训练预计需要1-2小时，请耐心等待..."

python src/trainer.py --train --wandb

if %errorlevel% equ 0 (
    call :log_success "模型训练完成！"
) else (
    call :log_error "模型训练失败！"
    exit /b 1
)
exit /b

REM 测试模型
:test_model
call :log_info "测试微调后的模型..."
python src/trainer.py --test
exit /b

REM 启动API服务
:start_api
call :log_info "启动API服务..."
call :log_info "服务将在 http://localhost:8000 启动"
call :log_info "API文档: http://localhost:8000/docs"

python src/api_server.py
exit /b

REM 显示帮助信息
:show_help
echo 佛山大学知识问答系统 - Windows启动脚本
echo.
echo 用法: %~nx0 [命令]
echo.
echo 命令:
echo   setup     - 初始化环境（检查模型、数据等）
echo   download  - 下载必要的模型文件
echo   data      - 处理数据和构建RAG向量库
echo   train     - 开始LoRA微调训练
echo   test      - 测试微调后的模型
echo   api       - 启动API服务
echo   all       - 完整流程（setup + train + api）
echo   help      - 显示此帮助信息
echo.
echo 示例:
echo   %~nx0 setup     # 初始化环境
echo   %~nx0 train     # 开始训练
echo   %~nx0 api       # 启动服务
echo.
exit /b

REM 完整流程
:run_all
call :log_info "执行完整流程..."

call :setup_environment
call :train_model
call :test_model

call :log_success "所有步骤完成！"

set /p choice="是否启动API服务？(y/n): "
if /i "!choice!"=="y" (
    call :start_api
)
exit /b

REM 环境设置
:setup_environment
call :log_info "设置环境..."

call :check_venv
call :activate_venv
call :check_cuda
call :check_models
call :check_data

call :log_success "环境设置完成！"
exit /b

REM 主函数
:main
if "%~1"=="" goto show_help
if "%~1"=="help" goto show_help
if "%~1"=="--help" goto show_help
if "%~1"=="-h" goto show_help

if "%~1"=="setup" (
    call :setup_environment
    goto :eof
)

if "%~1"=="download" (
    call :activate_venv
    call :download_models
    call :download_embedding_model
    goto :eof
)

if "%~1"=="data" (
    call :activate_venv
    call :check_data
    goto :eof
)

if "%~1"=="train" (
    call :activate_venv
    call :setup_environment
    call :train_model
    goto :eof
)

if "%~1"=="test" (
    call :activate_venv
    call :test_model
    goto :eof
)

if "%~1"=="api" (
    call :activate_venv
    call :start_api
    goto :eof
)

if "%~1"=="all" (
    call :run_all
    goto :eof
)

call :log_error "未知命令: %~1"
echo.
call :show_help
exit /b 1

REM 执行主函数
call :main %*